package top.kangyaocoding.tech.common;

import com.fasterxml.jackson.core.JsonParser;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.DeserializationContext;
import com.fasterxml.jackson.databind.JsonDeserializer;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ObjectNode;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.data.redis.serializer.RedisSerializer;

import java.io.IOException;

/**
 * 描述: Redis序列化Message
 *
 * @author K·Herbert
 * @since 2025-04-13 23:03
 */
public class MessageRedisSerializer implements RedisSerializer<Message> {

    private static final String MESSAGE_TYPE_FIELD = "messageType";
    private static final String TEXT_FIELD = "text";

    public enum MessageType {
        USER, ASSISTANT, SYSTEM;

        public static MessageType fromString(String type) {
            try {
                return MessageType.valueOf(type.toUpperCase());
            } catch (IllegalArgumentException e) {
                throw new RuntimeException("未知的消息类型: " + type);
            }
        }
    }

    private final ObjectMapper objectMapper;
    private final JsonDeserializer<Message> messageDeserializer;

    public MessageRedisSerializer(ObjectMapper objectMapper) {
        this.objectMapper = objectMapper;
        this.messageDeserializer = new JsonDeserializer<>() {
            @Override
            public Message deserialize(JsonParser jp, DeserializationContext ctx)
                    throws IOException {
                ObjectNode root = jp.readValueAsTree();
                if (!root.has(MESSAGE_TYPE_FIELD) || !root.has(TEXT_FIELD)) {
                    throw new RuntimeException("JSON 数据缺少必要字段: messageType 或 text");
                }
                String type = root.get(MESSAGE_TYPE_FIELD).asText();
                String text = root.get(TEXT_FIELD).asText();

                MessageType messageType = MessageType.fromString(type);
                return switch (messageType) {
                    case USER -> new UserMessage(text);
                    case ASSISTANT -> new AssistantMessage(text);
                    case SYSTEM -> new SystemMessage(text);
                };
            }
        };
    }

    @Override
    public byte[] serialize(Message message) {
        try {
            return objectMapper.writeValueAsBytes(message);
        } catch (JsonProcessingException e) {
            throw new RuntimeException("无法序列化消息: " + e.getMessage(), e);
        }
    }

    @Override
    public Message deserialize(byte[] bytes) {
        if (bytes == null || bytes.length == 0) {
            return null;
        }
        try {
            return messageDeserializer.deserialize(
                    objectMapper.getFactory().createParser(bytes),
                    objectMapper.getDeserializationContext()
            );
        } catch (IOException e) {
            throw new RuntimeException("无法反序列化消息: " + e.getMessage(), e);
        }
    }
}

